#!/usr/bin/python
import sys
sys.path.append('/home/a/unitree_sdk/unitree_legged_sdk-3.2/build') # Edit the path to "build" folder on your computer
import unitree

class Unitree_Robot_High():
    
    def __init__(self):
        import robot_interface_high_level as robot_interface_high
        self.unitree_robot = robot_interface_high.RobotInterface(unitree.HIGHLEVEL)
        self.robot_state = robot_interface_high.HighState()

        self.mode = 0
        self.footRaiseHeight = 0.0
        self.forwardSpeed = 0.0
        self.sideSpeed = 0.0
        self.rotateSpeed = 0.0
        self.bodyHeight = 0.0
        self.yaw = 0.0
        self.pitch = 0.0
        self.roll = 0.0

    def cmd_init(self):
        self.mode = 0
        self.footRaiseHeight = 0.0
        self.forwardSpeed = 0.0
        self.sideSpeed = 0.0
        self.rotateSpeed = 0.0
        self.bodyHeight = 0.0
        self.yaw = 0.0
        self.pitch = 0.0
        self.roll = 0.0
    
    def send_UDP(self):
        self.unitree_robot.UDPSend()

    def recv_UDP(self):
        self.unitree_robot.UDPRecv()
        
    def robot_control(self):
        self.unitree_robot.robotControl(self.mode,
                                    self.footRaiseHeight, self.bodyHeight, 
                                    self.roll, self.pitch, self.yaw,
                                    self.forwardSpeed, self.sideSpeed, self.rotateSpeed)
    
    
    def robot_walking(self, mode=2, forwardSpeed = 0.0, sideSpeed = 0.0, 
                      rotateSpeed = 0.0, bodyHeight = 0.0, footRaiseHeight = 0.0):
        # forwardSpeed: unit: m/s -1.0 ~ 1.0
    	# sideSpeed: unit: m/s -1.0 ~ 1.0
    	# rotateSpeed: unit: rad/s -1.0 ~ 1.0
    	# bodyHeight: unit: m, default: 0.28m
    	# footRaiseHeight
        self.cmd_init()
        self.footRaiseHeight = footRaiseHeight
        self.forwardSpeed = forwardSpeed + unitree.calibration_forwardspeed
        self.sideSpeed = sideSpeed + unitree.calibration_sidespeed
        self.rotateSpeed = rotateSpeed + unitree.calibration_rotatespeed
        self.bodyHeight = bodyHeight
        self.mode = mode
        self.recv_UDP()
        self.robot_control()
        self.send_UDP()
        return True
    

    def robot_pose(self, roll, pitch, yaw, bodyHeight):
        self.cmd_init()
        self.bodyHeight = bodyHeight
        self.yaw = yaw
        self.pitch = pitch
        self.roll = roll
        self.mode = 1
        self.recv_UDP()
        self.robot_control()
        self.send_UDP()
        return True 
    
    def getState(self):
        robot_state = self.unitree_robot.getState()
        return robot_state


#===============================================================================

class Unitree_Robot_Low():
    
    def __init__(self):
        import robot_interface_low_level as robot_interface_low
        self.unitree_robot = robot_interface_low.RobotInterface(unitree.LOWLEVEL)
        self.robot_state = robot_interface_low.LowState()

        # self.mode = 0
        self.motor = -1
        self.work = [0] * 12
        self.footRaiseHeight = [0.0] * 12
        self.forwardSpeed = [0.0] * 12
        self.sideSpeed = [0.0] * 12
        self.rotateSpeed = [0.0] * 12
        self.bodyHeight = [0.0] * 12
        self.yaw = [0.0] * 12
        self.pitch = [0.0] * 12
        self.roll = [0.0] * 12

    def cmd_init(self):
        # self.mode = 0
        self.motor = -1
        self.work = [0] * 12
        self.footRaiseHeight = [0.0] * 12
        self.forwardSpeed = [0.0] * 12
        self.sideSpeed = [0.0] * 12
        self.rotateSpeed = [0.0] * 12
        self.bodyHeight = [0.0] * 12
        self.yaw = [0.0] * 12
        self.pitch = [0.0] * 12
        self.roll = [0.0] * 12
    
    def send_UDP(self):
        self.unitree_robot.UDPSend()

    def recv_UDP(self):
        self.unitree_robot.UDPRecv()
        
    def robot_control(self):
        self.unitree_robot.robotControl(self.motor, self.work,
                                    # self.mode,
                                    self.q, self.dq, self.tau,
                                    self.Kp, self.Kd)
    
    
    def robot_work(self, motor=-1, work=[0]*12, 
                    angle=[0.0]*12, velocity=[0.0]*12, torque_output=[0.0]*12,
                    position_stiffness=[0.0]*12, velocity_stiffness=[0.0]*12):
        if motor==-1 and (not all(work)): return False
        self.cmd_init()
        self.motor = motor
        self.work = work
        # self.mode = mode
        self.q = angle
        self.dq = velocity
        self.tau = torque_output
        self.Kp = position_stiffness
        self.Kd = velocity_stiffness
        return True
    
    def getState(self):
        robot_state = self.unitree_robot.getState()
        return robot_state